# coding = utf-8

'''
基于loss的损失测试
'''

import sys

import click
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from pathlib2 import Path
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
from tqdm import tqdm

import utils.checkpoint as cp
from dataset import KiTSInstance
from dataset.transform import MedicalTransform
from loss import GeneralizedDiceLoss,CrossEntropyWeightLoss
from loss.util import class2one_hot
from network import     MUNet
from utils.metrics import Evaluator
from utils.vis import imshow

from utils import EarlyStopping

def test_loss():
    transform = MedicalTransform(output_size=(512, 512), roi_error_range=15, use_roi=False)

    # roi 为 None， 同时valid不做变换
    dataset = KiTSInstance("/datasets/DongbeiDaxue/chengkun_remove_only_tumor", stack_num=0, img_size=(512,512),
                           train_transform=None, valid_transform=None)

    net = MUNet(input_channel=dataset.img_channels, n_class=3)
    #net._initialize_weights()
    data = {'net': net}
    cp_file = Path(
        "/home/diaozhaoshuo/log/BeliefFunctionNN/chengkung/dongbeidaxue/munet/checkpoint_new_data_stack_0/cp_018.pth")
    cp.load_params(data, cp_file, device='cpu')
    net = net.cuda()



    subset = dataset.train_dataset
    case_slice_indices = dataset.train_case_slice_indices
    sampler = SequentialSampler(subset)
    data_loader = DataLoader(subset, batch_size=8, sampler=sampler,
                             num_workers=1, pin_memory=True)




    net = torch.nn.DataParallel(net, device_ids=[0, 1]).cuda()
    #net = net.cuda()
    net.eval()
    torch.set_grad_enabled(False)
    transform.eval()

    criterion1 = torch.nn.CrossEntropyLoss().cuda()
    criterion2 = CrossEntropyWeightLoss().cuda()

    for batch_idx, data in enumerate(data_loader):
        imgs, labels, idx, weight = data['image'].cuda(), data['label'].cuda(), data['index'], data["weight"].cuda()

        output = net(imgs)


        loss1 = criterion1(output, labels)
        loss2 = criterion2(output, labels, weight)
        print(loss1, loss2)






if __name__ == '__main__':
    test_loss()